#include <iostream>
#include <intrin.h>
#include <malloc.h>

extern "C" void avx2_add(float *a, float *b, float *c);

//判断是否支持avx2
bool is_avx2_supported()
{
    int cpuInfo[4];
    __cpuid(cpuInfo, 1);
    bool avx_supported = (cpuInfo[2]&(1<<28)) != 0;
    std::cout << "avx: " << avx_supported << std::endl;
    if(!avx_supported)
        return false;
    __cpuid(cpuInfo, 7);
    bool avx2_supported = (cpuInfo[1]&(1<<5)) != 0;
    std::cout << "avx2: " << avx2_supported << std::endl;
    return avx2_supported;
}

int main() 
{
    try{
        float *a = (float*)_aligned_malloc(8 * sizeof(float), 32);
        float *b = (float*)_aligned_malloc(8 * sizeof(float), 32);
        float *c = (float*)_aligned_malloc(8 * sizeof(float), 32);
    
        for (int i = 0; i < 8; i++) {
            a[i] = static_cast<float>(i + 1);
            b[i] = static_cast<float>(8 - i);
        }
    

        // 打印输入数组 a
        std::cout << "Array a:" << std::endl;
        for (int i = 0; i < 8; i++) {
            std::cout << "a[" << i << "] = " << a[i] << std::endl;
        }
    
        // 打印输入数组 b
        std::cout << "Array b:" << std::endl;
        for (int i = 0; i < 8; i++) {
            std::cout << "b[" << i << "] = " << b[i] << std::endl;
        }
    
        // 调用汇编函数
       if(is_avx2_supported())
       {
         std::cout << " is  support avx2!" << std::endl;
    
         avx2_add(a, b, c);
    
        // 输出结果
        std::cout << "Array c:" << std::endl;
        for (int i = 0; i < 8; i++) {
            std::cout << "c[" << i << "] = " << c[i] << std::endl;
        }
        
       }
       else
      {
        std::cout << " is not support avx2!" << std::endl;
       }
    
       _aligned_free(a);
       _aligned_free(b);
       _aligned_free(c);

    }catch (const std::exception& e) {
        std::cerr << "Exception caught: " << e.what() << std::endl;
        return 1;
    }

    return 0;
}


//version
; wht.asm
global fwht_8x8_row_pass_avx2
section .text

; fwht_8x8_row_pass_avx2 函数实现

fwht_8x8_row_pass_avx2:
    ; 保存寄存器状态
    push rbp
    mov rbp, rsp
    push rbx
    push r12
    push r13
    push r14
    push r15

    ; 获取参数（Windows调用约定）
    mov r12, rcx ; src
    mov r13, rdx ; dst
    mov r14, r8  ; row
    mov r15d, r9d ; dpit

    ; 循环处理每一行
.loop_rows:
    ; 加载数据, 
    vmovdqa xmm0, [r12]           ; xmm0 = [sptr[0], sptr[1], sptr[2], sptr[3], sptr[4], sptr[5], sptr[6], sptr[7]]
    vpshufd xmm1, xmm0, 0b11100100 ; xmm1 = [sptr[4], sptr[5], sptr[6], sptr[7], sptr[0], sptr[1], sptr[2], sptr[3]]
    
    ; 第一步：计算 bdata
    vpaddw xmm2, xmm0, xmm1        ; xmm2 = [sptr[0]+sptr[4], sptr[1]+sptr[5], sptr[2]+sptr[6], sptr[3]+sptr[7], sptr[4]+sptr[0], sptr[5]+sptr[1], sptr[6]+sptr[2], sptr[7]+sptr[3]]
    vpsubw xmm3, xmm0, xmm1        ; xmm3 = [sptr[0]-sptr[4], sptr[1]-sptr[5], sptr[2]-sptr[6], sptr[3]-sptr[7], sptr[4]-sptr[0], sptr[5]-sptr[1], sptr[6]-sptr[2], sptr[7]-sptr[3]]

    ; 重新排列 xmm2 和 xmm3 以进行下一步计算
    vpshufd xmm4, xmm2, 0b10001101 ; xmm4 = [bdata[0], bdata[2], bdata[1], bdata[3], bdata[4], bdata[6], bdata[5], bdata[7]]
    vpshufd xmm5, xmm3, 0b10001101 ; xmm5 = [bdata[0], bdata[2], bdata[1], bdata[3], bdata[4], bdata[6], bdata[5], bdata[7]]

    ; 第二步：计算 adata
    vpaddw xmm6, xmm4, xmm5        ; xmm6 = [bdata[0]+bdata[2], bdata[1]+bdata[3], bdata[0]+bdata[2], bdata[1]+bdata[3], bdata[4]+bdata[6], bdata[5]+bdata[7], bdata[4]+bdata[6], bdata[5]+bdata[7]]
    vpsubw xmm7, xmm4, xmm5        ; xmm7 = [bdata[0]-bdata[2], bdata[1]-bdata[3], bdata[0]-bdata[2], bdata[1]-bdata[3], bdata[4]-bdata[6], bdata[5]-bdata[7], bdata[4]-bdata[6], bdata[5]-bdata[7]]

    ; 重新排列 xmm6 和 xmm7 以进行下一步计算
    vpshufd xmm8, xmm6, 0b11001001 ; xmm8 = [adata[0], adata[1], adata[2], adata[3], adata[4], adata[5], adata[6], adata[7]]
    vpshufd xmm9, xmm7, 0b11001001 ; xmm9 = [adata[0], adata[1], adata[2], adata[3], adata[4], adata[5], adata[6], adata[7]]

    ; 第三步：计算最终结果
    vpaddw xmm10, xmm8, xmm9       ; xmm10 = [adata[0]+adata[1], adata[0]-adata[1], adata[2]+adata[3], adata[2]-adata[3], adata[4]+adata[5], adata[4]-adata[5], adata[6]+adata[7], adata[6]-adata[7]]

    ; 重新排列最终结果以匹配 C++ 存储逻辑
    vpshufd xmm11, xmm10, 0b00000000 ; 这里根据实际情况调整重排控制码
    
    ; 存储结果
     mov rbx, r13
     vmovdqu [rbx], xmm11

    ; 更新指针
    add r12, 16  ;8 * 2 bytes
    lea rcx, [r15d * 8] 
    add r13, rcx
   

    ; 循环计数
    dec r14
    jnz .loop_rows

    ; 恢复寄存器状态
    pop r15
    pop r14
    pop r13
    pop r12
    pop rbx
    pop rbp

    ret